/*
 * SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES.
 * All rights reserved. SPDX-License-Identifier: LicenseRef-NvidiaProprietary
 *
 * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
 * property and proprietary rights in and to this material, related
 * documentation and any modifications thereto. Any use, reproduction,
 * disclosure or distribution of this material and related documentation
 * without an express license agreement from NVIDIA CORPORATION or
 * its affiliates is strictly prohibited.
 */

#ifndef TRT_MULTISCALE_DEFORMABLE_ATTN_PLUGIN_H
#define TRT_MULTISCALE_DEFORMABLE_ATTN_PLUGIN_H

// For loadLibrary
#ifdef _MSC_VER
// Needed so that the max/min definitions in windows.h do not conflict with
// std::max/min.
#define NOMINMAX
#include <windows.h>
#undef NOMINMAX
#else
#include <dlfcn.h>
#endif

#include <iostream>
#include <memory>
#include <sstream>
#include <string>
#include <vector>

#include "NvInfer.h"
#include "NvInferPlugin.h"
#include "NvInferVersion.h"
#include "msda.h"

#if NV_TENSORRT_MAJOR >= 10

#if NV_TENSORRT_MAJOR > 7
#define PLUGIN_NOEXCEPT noexcept
#else
#define PLUGIN_NOEXCEPT
#endif

using namespace nvinfer1::plugin;

namespace nvinfer1 {
namespace plugin {

namespace {
static char const* MSDA_VERSION{"1"};
static char const* MSDA_NAME{"MSDA"};
static constexpr int MSDA_MAX_SCALES{32};
}  // namespace

class MSDAPluginV3 : public IPluginV3,
                     public IPluginV3OneCore,
                     public IPluginV3OneBuild,
                     public IPluginV3OneRuntime {
  struct Params;

 public:
  MSDAPluginV3() = default;

  MSDAPluginV3(const Params& params);

  MSDAPluginV3(const int2* spatial_shape, const float* scale_value,
               float scale_offset, float scale_weight, float scale_out,
               Layout source_layout);

  IPluginCapability* getCapabilityInterface(PluginCapabilityType type)
      PLUGIN_NOEXCEPT override;

  IPluginV3* clone() PLUGIN_NOEXCEPT override;
  int getOutputShapes(DimsExprs const* inputs, int32_t nbInputs,
                      DimsExprs const* shapeInputs, int32_t nbShapeInputs,
                      DimsExprs* outputs, int32_t nbOutputs,
                      IExprBuilder& exprBuilder) PLUGIN_NOEXCEPT override;
  bool supportsFormatCombination(int32_t pos,
                                 DynamicPluginTensorDesc const* inOut,
                                 int32_t nbInputs,
                                 int32_t nbOutputs) PLUGIN_NOEXCEPT override;
  int32_t configurePlugin(DynamicPluginTensorDesc const* in, int32_t nbInputs,
                          DynamicPluginTensorDesc const* out,
                          int32_t nbOutputs) PLUGIN_NOEXCEPT override;
  size_t getWorkspaceSize(DynamicPluginTensorDesc const* inputs,
                          int32_t nbInputs,
                          DynamicPluginTensorDesc const* outputs,
                          int32_t nbOutputs) const PLUGIN_NOEXCEPT override;
  int32_t enqueue(PluginTensorDesc const* inputDesc,
                  PluginTensorDesc const* outputDesc, void const* const* inputs,
                  void* const* outputs, void* workspace,
                  cudaStream_t stream) PLUGIN_NOEXCEPT override;
  int32_t onShapeChange(PluginTensorDesc const* in, int32_t nbInputs,
                        PluginTensorDesc const* out,
                        int32_t nbOutputs) PLUGIN_NOEXCEPT override;
  IPluginV3* attachToContext(IPluginResourceContext* context)
      PLUGIN_NOEXCEPT override;

  // IPluginV2Ext Methods
  int32_t getOutputDataTypes(DataType* outputTypes, int32_t nbOutputs,
                             DataType const* inputTypes,
                             int32_t nbInputs) const PLUGIN_NOEXCEPT override;

  // IPluginV2 Methods
  char const* getPluginName() const PLUGIN_NOEXCEPT override;
  char const* getPluginVersion() const PLUGIN_NOEXCEPT override;
  int32_t getNbOutputs() const PLUGIN_NOEXCEPT override;

  PluginFieldCollection const* getFieldsToSerialize() PLUGIN_NOEXCEPT override;
  char const* getPluginNamespace() const PLUGIN_NOEXCEPT override;

 private:
  std::string mNamespace;
  struct Params {
    int2 spatial_shape[MSDA_MAX_SCALES];
    float scale_value[MSDA_MAX_SCALES];
    float scale_offset;
    float scale_weight;
    float scale_out;
    Layout source_layout;
    PluginField data_serialize[6];
    PluginFieldCollection fields_serialize;
    const PluginFieldCollection* getFields() {
      data_serialize[0] = {"spatial_shape", spatial_shape,
                           PluginFieldType::kINT32, MSDA_MAX_SCALES * 2};
      data_serialize[1] = {"scale_value", scale_value,
                           PluginFieldType::kFLOAT32, MSDA_MAX_SCALES};
      data_serialize[2] = {"scale_offset", &scale_offset,
                           PluginFieldType::kFLOAT32, 1};
      data_serialize[3] = {"scale_weight", &scale_weight,
                           PluginFieldType::kFLOAT32, 1};
      data_serialize[4] = {"scale_out", &scale_out, PluginFieldType::kFLOAT32,
                           1};
      data_serialize[5] = {"feature_layout", &source_layout,
                           PluginFieldType::kINT32, 1};
      fields_serialize.nbFields = 6;
      fields_serialize.fields = data_serialize;
      return &fields_serialize;
    };
  } params;

#if NV_TENSORRT_MAJOR < 8
  using nvinfer1::IPluginV2DynamicExt::canBroadcastInputAcrossBatch;
  using nvinfer1::IPluginV2DynamicExt::configurePlugin;
  using nvinfer1::IPluginV2DynamicExt::enqueue;
  using nvinfer1::IPluginV2DynamicExt::getOutputDimensions;
  using nvinfer1::IPluginV2DynamicExt::getWorkspaceSize;
  using nvinfer1::IPluginV2DynamicExt::isOutputBroadcastAcrossBatch;
  using nvinfer1::IPluginV2DynamicExt::supportsFormat;
#endif
};

class MSDAPluginV3Creator : public nvinfer1::IPluginCreatorV3One {
 public:
  MSDAPluginV3Creator();
  char const* getPluginName() const PLUGIN_NOEXCEPT override;
  char const* getPluginVersion() const PLUGIN_NOEXCEPT override;
  PluginFieldCollection const* getFieldNames() PLUGIN_NOEXCEPT override;
  IPluginV3* createPlugin(char const* name, PluginFieldCollection const* fc,
                          TensorRTPhase phase) PLUGIN_NOEXCEPT override;
  char const* getPluginNamespace() const PLUGIN_NOEXCEPT override;

 private:
  PluginFieldCollection mFC;
  std::vector<nvinfer1::PluginField> mPluginAttributes;
  std::string mNamespace = "";
};

}  // namespace plugin
}  // namespace nvinfer1

#endif
#endif
